#!/usr/bin/env python3
# coding=utf-8
# @author Loganli
'''
Author: Logan.Li
Gitee: https://gitee.com/attacker
email: admin@attacker.club
Date: 2025-05-05 22:28:07
LastEditTime: 2025-05-05 22:28:07
Description: AWS EC2 Zabbix Agent 批量部署工具 (基于paramiko+boto3)
'''

import boto3
import paramiko
import argparse
import os
import sys
import concurrent.futures
from botocore.exceptions import ClientError

# 日志颜色
RED = '\033[31m'
GREEN = '\033[32m'
YELLOW = '\033[33m'
BLUE = '\033[34m'
BOLD = '\033[1m'
CYAN = '\033[36m'
RESET = '\033[0m'

def log_info(message):
    print(f"{BLUE}[信息]{RESET} {message}")

def log_success(message):
    print(f"{GREEN}[成功]{RESET} {message}")

def log_warning(message):
    print(f"{YELLOW}[警告]{RESET} {message}")

def log_error(message):
    print(f"{RED}[错误]{RESET} {message}")

class ZabbixDeployer:
    def __init__(self, region, profile=None):
        """初始化AWS连接"""
        try:
            log_info(f"使用{'默认' if not profile else profile}凭证连接AWS区域 {region}")
            if profile:
                self.session = boto3.Session(profile_name=profile, region_name=region)
            else:
                self.session = boto3.Session(region_name=region)
            
            self.ec2_client = self.session.client('ec2')
            self.ec2_resource = self.session.resource('ec2')
            log_success("AWS连接初始化成功")
        except Exception as e:
            log_error(f"AWS连接初始化失败: {str(e)}")
            sys.exit(1)
    
    def get_instances(self, filters=None):
        """获取EC2实例列表"""
        if filters is None:
            filters = []
            
        try:
            # 添加默认过滤器：只获取运行中的实例
            filters.append({'Name': 'instance-state-name', 'Values': ['running']})
            
            instances = self.ec2_resource.instances.filter(Filters=filters)
            instance_list = []
            
            for instance in instances:
                name = ''
                for tag in instance.tags or []:
                    if tag['Key'] == 'Name':
                        name = tag['Value']
                        break
                
                # 获取实例信息
                instance_info = {
                    'id': instance.id,
                    'name': name,
                    'private_ip': instance.private_ip_address,
                    'public_ip': instance.public_ip_address,
                    'key_name': instance.key_name,
                    'platform': instance.platform,
                    'image_id': instance.image_id
                }
                
                # 确定SSH用户名（根据AMI类型）
                if instance.platform == 'windows':
                    instance_info['ssh_user'] = 'Administrator'
                else:
                    # 根据AMI ID判断Linux发行版类型
                    ami_id = instance.image_id
                    try:
                        ami = self.ec2_client.describe_images(ImageIds=[ami_id])
                        if ami['Images']:
                            description = ami['Images'][0].get('Description', '').lower()
                            if 'ubuntu' in description:
                                instance_info['ssh_user'] = 'ubuntu'
                            elif 'debian' in description:
                                instance_info['ssh_user'] = 'admin'
                            elif 'centos' in description or 'rhel' in description or 'amazon' in description:
                                instance_info['ssh_user'] = 'ec2-user'
                            else:
                                instance_info['ssh_user'] = 'ec2-user'  # 默认
                    except:
                        instance_info['ssh_user'] = 'ec2-user'  # 出错时使用默认值
                
                instance_list.append(instance_info)
            
            log_success(f"成功获取 {len(instance_list)} 个运行中的EC2实例")
            return instance_list
            
        except Exception as e:
            log_error(f"获取EC2实例列表失败: {str(e)}")
            return []

    def deploy_zabbix_agent(self, instance, ssh_key_path, zabbix_server, timeout=30, port=22):
        """部署Zabbix Agent到指定实例"""
        # 跳过Windows实例
        if instance.get('platform') == 'windows':
            log_warning(f"跳过Windows实例 {instance['name']} ({instance['id']})")
            return False
        
        # 确定连接IP
        host_ip = instance['private_ip'] or instance['public_ip']
        if not host_ip:
            log_error(f"实例 {instance['name']} ({instance['id']}) 没有可用IP地址")
            return False
        
        # 确定主机名
        hostname = instance['name'] or instance['id']
        
        # 确定SSH用户名
        ssh_user = instance.get('ssh_user', 'ec2-user')
        
        # 检查SSH密钥
        if not os.path.exists(ssh_key_path):
            log_error(f"SSH密钥文件不存在: {ssh_key_path}")
            return False
        
        # 检查SSH密钥权限
        key_stat = os.stat(ssh_key_path)
        if key_stat.st_mode & 0o077:
            log_warning(f"SSH密钥权限过于开放，正在修复: {ssh_key_path}")
            try:
                os.chmod(ssh_key_path, 0o600)
            except Exception as e:
                log_error(f"无法修改SSH密钥权限: {str(e)}")
                return False
        
        log_info(f"正在连接 {hostname} ({host_ip}) 使用用户 {ssh_user}...")
        
        try:
            # 创建SSH客户端
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # 加载私钥
            try:
                key = paramiko.RSAKey.from_private_key_file(ssh_key_path)
                ssh.connect(hostname=host_ip, username=ssh_user, pkey=key, timeout=timeout, port=port)
            except paramiko.ssh_exception.PasswordRequiredException:
                log_error(f"SSH密钥需要密码，不支持: {ssh_key_path}")
                return False
            except Exception as e:
                log_error(f"SSH连接失败 {hostname} ({host_ip}): {str(e)}")
                return False
            
            log_success(f"成功连接到 {hostname} ({host_ip})")
            
            # 执行Zabbix Agent安装命令
            install_cmd = f"curl -s https://gitee.com/attacker/deploy-zabbix/raw/master/agent/zabbix-agent-universal.sh | sudo bash -s {zabbix_server} {hostname}"
            log_info(f"正在执行安装命令: {install_cmd}")
            
            stdin, stdout, stderr = ssh.exec_command(install_cmd)
            exit_code = stdout.channel.recv_exit_status()
            
            # 获取命令输出
            stdout_str = stdout.read().decode('utf-8', errors='ignore')
            stderr_str = stderr.read().decode('utf-8', errors='ignore')
            
            if exit_code != 0:
                log_error(f"安装失败 {hostname} ({host_ip}), 退出码: {exit_code}")
                log_error(f"错误输出: {stderr_str}")
                return False
            
            # 检查安装结果
            check_cmd = "sudo systemctl status zabbix-agent || sudo service zabbix-agent status"
            stdin, stdout, stderr = ssh.exec_command(check_cmd)
            status_output = stdout.read().decode('utf-8', errors='ignore')
            
            if "active (running)" in status_output or "is running" in status_output:
                log_success(f"Zabbix Agent 在 {hostname} ({host_ip}) 安装并运行成功")
                return True
            else:
                log_warning(f"Zabbix Agent 在 {hostname} ({host_ip}) 安装完成，但服务可能未运行")
                return False
                
        except Exception as e:
            log_error(f"部署过程出错 {hostname} ({host_ip}): {str(e)}")
            return False
        finally:
            try:
                ssh.close()
            except:
                pass

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='AWS EC2 Zabbix Agent 批量部署工具')
    parser.add_argument('--zabbix-server', required=True, help='Zabbix Server IP地址')
    parser.add_argument('--ssh-key', required=True, help='SSH私钥文件路径')
    parser.add_argument('--region', default='ap-northeast-1', help='AWS区域 (默认: ap-northeast-1)')
    parser.add_argument('--profile', help='AWS配置文件名称')
    parser.add_argument('--filter', help='实例名称过滤条件')
    parser.add_argument('--port', type=int, default=22, help='SSH端口 (默认: 22)')
    parser.add_argument('--timeout', type=int, default=30, help='SSH连接超时时间 (默认: 30秒)')
    parser.add_argument('--max-workers', type=int, default=5, help='最大并发数 (默认: 5)')
    
    args = parser.parse_args()
    
    # 初始化部署器
    deployer = ZabbixDeployer(args.region, args.profile)
    
    # 获取实例列表
    filters = []
    if args.filter:
        filters.append({
            'Name': 'tag:Name',
            'Values': [f'*{args.filter}*']
        })
    
    instances = deployer.get_instances(filters)
    
    if not instances:
        log_error("没有找到符合条件的EC2实例")
        sys.exit(1)
    
    # 显示找到的实例
    log_info("找到以下EC2实例:")
    
    # 美化表格输出
    fmt = "{:<3} {:<20} {:<20} {:<15} {:<15}"
    print(f"\n{BOLD}{fmt.format('ID', '实例名称', '实例ID', 'IP地址', '用户')}{RESET}")
    print("─" * 80)  # 使用更美观的分隔线
    
    for i, instance in enumerate(instances, 1):
        ip = instance['private_ip'] or instance['public_ip'] or 'N/A'
        name = instance['name'] if instance['name'] else 'N/A'
        print(fmt.format(
            f"{i}.", 
            name[:18] + '..' if len(name) > 20 else name, 
            instance['id'], 
            ip, 
            instance.get('ssh_user', 'ec2-user')
        ))
    print("─" * 80)
    
    # 确认是否继续
    confirm = input(f"\n{BOLD}是否继续部署? (y/n):{RESET} ").strip().lower()
    if confirm != 'y':
        log_info("用户取消部署")
        sys.exit(0)
    
    # 并发部署
    success_count = 0
    failed_count = 0
    
    print(f"\n{BOLD}{CYAN}▶ 开始批量部署 Zabbix Agent...{RESET}")
    print("─" * 80)
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.max_workers) as executor:
        future_to_instance = {
            executor.submit(
                deployer.deploy_zabbix_agent, 
                instance, 
                args.ssh_key, 
                args.zabbix_server,
                args.timeout,
                args.port
            ): instance for instance in instances
        }
        
        for future in concurrent.futures.as_completed(future_to_instance):
            instance = future_to_instance[future]
            try:
                result = future.result()
                if result:
                    success_count += 1
                else:
                    failed_count += 1
            except Exception as e:
                log_error(f"部署异常 {instance['name']} ({instance['id']}): {str(e)}")
                failed_count += 1
    
    # 显示部署结果
    print("\n" + "═" * 80)
    log_info(f"部署完成: {GREEN}成功 {success_count} 个{RESET}, {RED}失败 {failed_count} 个{RESET}")
    print("═" * 80)

if __name__ == "__main__":
    main()
